//===----------------------------------------------------------------------===//
//
//                         BusTub
//
// lock_manager.cpp
//
// Identification: src/concurrency/lock_manager.cpp
//
// Copyright (c) 2015-2019, Carnegie Mellon University Database Group
//
//===----------------------------------------------------------------------===//

#include "concurrency/lock_manager.h"

#include "common/config.h"
#include "concurrency/transaction.h"
#include "concurrency/transaction_manager.h"

namespace bustub {

/**
 * @param txn Transaction
 * @param mode The lock mode attempt on the row
 * In lock row, we should only allow shared lock or exclusive lock, the intention lock will fail
 */
void LockManager::CheckLocksAllowedOnRows(Transaction *txn, LockMode mode) {
  if (mode != LockMode::SHARED && mode != LockMode::EXCLUSIVE) {
    txn->SetState(TransactionState::ABORTED);
    throw TransactionAbortException(txn->GetTransactionId(), AbortReason::ATTEMPTED_INTENTION_LOCK_ON_ROW);
  }
}

void LockManager::HoldTableLockBeforeApplyRowLock(Transaction *txn, LockMode mode, table_oid_t table_id) {
  if (mode == LockMode::SHARED) {
    auto any_shared_lock_hold = txn->GetSharedIntentionExclusiveTableLockSet()->count(table_id) != 0 ||
                                txn->GetSharedTableLockSet()->count(table_id) != 0 ||
                                txn->GetIntentionSharedTableLockSet()->count(table_id) != 0 ||
                                txn->GetIntentionExclusiveTableLockSet()->count(table_id) != 0 ||
                                txn->GetExclusiveTableLockSet()->count(table_id) != 0;
    if (any_shared_lock_hold) {
      return;
    }
  } else if (mode == LockMode::EXCLUSIVE) {
    auto any_exclusive_lock_hold = txn->GetSharedIntentionExclusiveTableLockSet()->count(table_id) != 0 ||
                                   txn->GetIntentionExclusiveTableLockSet()->count(table_id) != 0 ||
                                   txn->GetExclusiveTableLockSet()->count(table_id) != 0;
    if (any_exclusive_lock_hold) {
      return;
    }
  }
  txn->SetState(TransactionState::ABORTED);
  throw TransactionAbortException(txn->GetTransactionId(), AbortReason::TABLE_LOCK_NOT_PRESENT);
}

void LockManager::TransactionLevelCheck(Transaction *txn, LockMode mode) {
  if (txn->GetState() == TransactionState::COMMITTED) {
    txn->SetState(TransactionState::GROWING);
  }
  switch (txn->GetIsolationLevel()) {
    case IsolationLevel::REPEATABLE_READ:
      // Locks should be granted only in growing state
      if (txn->GetState() != TransactionState::GROWING) {
        txn->SetState(TransactionState::ABORTED);
        throw TransactionAbortException(txn->GetTransactionId(), AbortReason::LOCK_ON_SHRINKING);
      }
      return;
    case IsolationLevel::READ_COMMITTED:
      // All locks allowed in read_committed, only Is, S locks allowed in shrinking
      if (txn->GetState() != TransactionState::GROWING) {
        if (mode != LockMode::SHARED && mode != LockMode::INTENTION_SHARED) {
          txn->SetState(TransactionState::ABORTED);
          throw TransactionAbortException(txn->GetTransactionId(), AbortReason::LOCK_ON_SHRINKING);
        }
      }
      return;
    case IsolationLevel::READ_UNCOMMITTED:
      if (mode != LockMode::INTENTION_EXCLUSIVE && mode != LockMode::EXCLUSIVE) {
        txn->SetState(TransactionState::ABORTED);
        throw TransactionAbortException(txn->GetTransactionId(), AbortReason::LOCK_SHARED_ON_READ_UNCOMMITTED);
      }
      if (txn->GetState() != TransactionState::GROWING) {
        txn->SetState(TransactionState::ABORTED);
        throw TransactionAbortException(txn->GetTransactionId(), AbortReason::LOCK_ON_SHRINKING);
      }
      return;
    default:
      // No operation, just for passing the C-lang format test
      (void)txn;
  }
}

void LockManager::IsValidUpgradeRequest(Transaction *txn, LockMode origin_mode, LockMode upgrade_mode) {
  switch (origin_mode) {
    case LockMode::INTENTION_SHARED:
      if (upgrade_mode == LockMode::INTENTION_SHARED) {
        txn->SetState(TransactionState::ABORTED);
        throw TransactionAbortException(txn->GetTransactionId(), AbortReason::INCOMPATIBLE_UPGRADE);
      }
      return;
    case LockMode::SHARED:
    case LockMode::INTENTION_EXCLUSIVE:
      if (upgrade_mode != LockMode::SHARED_INTENTION_EXCLUSIVE && upgrade_mode != LockMode::EXCLUSIVE) {
        txn->SetState(TransactionState::ABORTED);
        throw TransactionAbortException(txn->GetTransactionId(), AbortReason::INCOMPATIBLE_UPGRADE);
      }
      return;
    case LockMode::SHARED_INTENTION_EXCLUSIVE:
      if (upgrade_mode != LockMode::EXCLUSIVE) {
        txn->SetState(TransactionState::ABORTED);
        throw TransactionAbortException(txn->GetTransactionId(), AbortReason::INCOMPATIBLE_UPGRADE);
      }
      return;
    default:
      txn->SetState(TransactionState::ABORTED);
      throw TransactionAbortException(txn->GetTransactionId(), AbortReason::INCOMPATIBLE_UPGRADE);
  }
}

void LockManager::HoldNoRowLocksBeforeUnlockTable(Transaction *txn, table_oid_t table_id) {
  auto no_row_lock =
      (txn->GetSharedRowLockSet()->count(table_id) == 0 || txn->GetSharedRowLockSet()->at(table_id).empty()) &&
      (txn->GetExclusiveRowLockSet()->count(table_id) == 0 || txn->GetExclusiveRowLockSet()->at(table_id).empty());
  if (!no_row_lock) {
    txn->SetState(TransactionState::ABORTED);
    throw TransactionAbortException(txn->GetTransactionId(), AbortReason::TABLE_UNLOCKED_BEFORE_UNLOCKING_ROWS);
  }
}

void LockManager::ChangeTransactionStateOnUnlock(Transaction *txn, LockMode mode) {
  switch (txn->GetIsolationLevel()) {
    case IsolationLevel::REPEATABLE_READ:
      if (mode == LockMode::SHARED || mode == LockMode::EXCLUSIVE) {
        txn->SetState(TransactionState::SHRINKING);
      }
      return;
    case IsolationLevel::READ_COMMITTED:
    case IsolationLevel::READ_UNCOMMITTED:
      if (mode == LockMode::EXCLUSIVE) {
        txn->SetState(TransactionState::SHRINKING);
      }
      return;
  }
}

auto LockManager::GetExistingLockRequestFromGrantQueue(txn_id_t txn_id, std::shared_ptr<LockRequestQueue> &queue_ptr)
    -> std::list<LockRequest *>::iterator {
  auto it = queue_ptr->granted_queue_.begin();
  while (it != queue_ptr->granted_queue_.end()) {
    if ((*it)->txn_id_ == txn_id) {
      break;
    }
    it++;
  }
  return it;
}

void LockManager::MoveGrantedRequestToGrantedQueue(std::shared_ptr<LockRequestQueue> &queue_ptr, LockRequest *request) {
  queue_ptr->request_queue_.remove(request);
  queue_ptr->granted_queue_.push_back(request);
}

void LockManager::CreateEntryInTableLockMap(Transaction *txn, table_oid_t table_id, LockMode lock_mode) {
  auto queue_ptr = std::make_shared<LockRequestQueue>();
  auto *request = new LockRequest{txn->GetTransactionId(), lock_mode, table_id};
  request->granted_ = true;
  queue_ptr->granted_queue_.push_back(request);
  table_lock_map_[table_id] = queue_ptr;
}

void LockManager::CreateEntryInTableRowLockMap(Transaction *txn, table_oid_t table_id, RID rid, LockMode lock_mode) {
  auto queue_ptr = std::make_shared<LockRequestQueue>();
  auto *request = new LockRequest{txn->GetTransactionId(), lock_mode, table_id, rid};
  request->granted_ = true;
  queue_ptr->granted_queue_.push_back(request);
  row_lock_map_[rid] = queue_ptr;
}

auto LockManager::IsTransactionHoldTableLock(Transaction *txn, table_oid_t oid, LockMode *mode) -> bool {
  if (txn->GetIntentionSharedTableLockSet()->find(oid) != txn->GetIntentionSharedTableLockSet()->end()) {
    *mode = LockMode::INTENTION_SHARED;
    return true;
  }
  if (txn->GetIntentionExclusiveTableLockSet()->find(oid) != txn->GetIntentionExclusiveTableLockSet()->end()) {
    *mode = LockMode::INTENTION_EXCLUSIVE;
    return true;
  }
  if (txn->GetSharedIntentionExclusiveTableLockSet()->find(oid) !=
      txn->GetSharedIntentionExclusiveTableLockSet()->end()) {
    *mode = LockMode::SHARED_INTENTION_EXCLUSIVE;
    return true;
  }
  if (txn->GetSharedTableLockSet()->find(oid) != txn->GetSharedTableLockSet()->end()) {
    *mode = LockMode::SHARED;
    return true;
  }
  if (txn->GetExclusiveTableLockSet()->find(oid) != txn->GetExclusiveTableLockSet()->end()) {
    *mode = LockMode::EXCLUSIVE;
    return true;
  }
  return false;
}

auto LockManager::IsTransactionHoldTableRowLock(Transaction *txn, table_oid_t oid, RID rid, LockMode *mode) -> bool {
  if (txn->GetSharedRowLockSet()->find(oid) != txn->GetSharedRowLockSet()->end() &&
      txn->GetSharedRowLockSet()->at(oid).count(rid) != 0) {
    *mode = LockMode::SHARED;
    return true;
  }
  if (txn->GetExclusiveRowLockSet()->find(oid) != txn->GetExclusiveRowLockSet()->end() &&
      txn->GetExclusiveRowLockSet()->at(oid).count(rid) != 0) {
    *mode = LockMode::EXCLUSIVE;
    return true;
  }
  return false;
}

void LockManager::GetLockStatesFromQueue(std::unordered_set<LockMode> &lock_modes,
                                         std::shared_ptr<LockRequestQueue> &queue_ptr, bool mode, txn_id_t txn_id) {
  auto iter = queue_ptr->granted_queue_.begin();
  while (iter != queue_ptr->granted_queue_.end()) {
    if (mode && (*iter)->txn_id_ == txn_id) {
      iter++;
      continue;
    }
    lock_modes.insert((*iter)->lock_mode_);
    iter++;
  }
}

auto LockManager::CompatibleWithLockModeSet(LockMode mode, std::unordered_set<LockMode> &table_lock_set) -> bool {
  auto it = table_lock_set.begin();
  while (it != table_lock_set.end()) {
    if (!Compatible(*it, mode)) {
      return false;
    }
    it++;
  }
  return true;
}

auto LockManager::Compatible(LockMode mode, LockMode other_mode) -> bool {
  switch (mode) {
    case LockMode::EXCLUSIVE:
      return false;
    case LockMode::INTENTION_SHARED:
      return other_mode != LockMode::EXCLUSIVE;
    case LockMode::INTENTION_EXCLUSIVE:
      return other_mode == LockMode::INTENTION_SHARED || other_mode == LockMode::INTENTION_EXCLUSIVE;
    case LockMode::SHARED:
      return other_mode == LockMode::INTENTION_SHARED || other_mode == LockMode::SHARED;
    case LockMode::SHARED_INTENTION_EXCLUSIVE:
      return other_mode == LockMode::INTENTION_SHARED;
  }
  return false;
}

auto LockManager::IsLockModeCompatible(LockMode mode, std::shared_ptr<LockRequestQueue> &lock_request_queue) -> bool {
  auto it = lock_request_queue->granted_queue_.begin();
  for (; it != lock_request_queue->granted_queue_.end(); it++) {
    if (!Compatible(mode, (*it)->lock_mode_)) {
      return false;
    }
  }
  return true;
}

auto LockManager::BlockedUtilRequestGranted(Transaction *txn, std::unique_lock<std::mutex> &unique_lock,
                                            std::condition_variable &condition, LockRequest *request,
                                            std::shared_ptr<LockRequestQueue> &queue_ptr) -> bool {
  while (true) {
    condition.wait(unique_lock,
                   [request, txn]() { return request->granted_ || txn->GetState() == TransactionState::ABORTED; });
    if (txn->GetState() != TransactionState::ABORTED && queue_ptr->request_queue_.front() != request) {
      condition.notify_all();
      continue;
    }
    if (txn->GetState() == TransactionState::ABORTED || request->granted_) {
      break;
    }
  }
  return txn->GetState() != TransactionState::ABORTED;
}

inline auto LockManager::FindLockRequestWithGivenId(std::shared_ptr<LockRequestQueue> &queue_ptr, txn_id_t txn_id)
    -> std::list<LockRequest *>::iterator {
  auto it = queue_ptr->request_queue_.begin();
  while (it != queue_ptr->request_queue_.end()) {
    if ((*it)->granted_) {
      if ((*it)->txn_id_ == txn_id) {
        return it;
      }
      it++;
    } else {
      break;
    }
  }
  return queue_ptr->request_queue_.end();
}

inline void LockManager::DeleteLockRecordFromTransaction(LockMode mode, Transaction *txn, const table_oid_t &table_id) {
  switch (mode) {
    case LockMode::INTENTION_SHARED:
      txn->GetIntentionSharedTableLockSet()->erase(table_id);
      break;
    case LockMode::INTENTION_EXCLUSIVE:
      txn->GetIntentionExclusiveTableLockSet()->erase(table_id);
      break;
    case LockMode::SHARED_INTENTION_EXCLUSIVE:
      txn->GetSharedIntentionExclusiveTableLockSet()->erase(table_id);
      break;
    case LockMode::SHARED:
      txn->GetSharedTableLockSet()->erase(table_id);
      break;
    case LockMode::EXCLUSIVE:
      txn->GetExclusiveTableLockSet()->erase(table_id);
      break;
  }
}

inline void LockManager::InsertLockRecordFromTransaction(LockMode mode, Transaction *txn, const table_oid_t &table_id) {
  switch (mode) {
    case LockMode::INTENTION_SHARED:
      txn->GetIntentionSharedTableLockSet()->insert(table_id);
      break;
    case LockMode::INTENTION_EXCLUSIVE:
      txn->GetIntentionExclusiveTableLockSet()->insert(table_id);
      break;
    case LockMode::SHARED_INTENTION_EXCLUSIVE:
      txn->GetSharedIntentionExclusiveTableLockSet()->insert(table_id);
      break;
    case LockMode::SHARED:
      txn->GetSharedTableLockSet()->insert(table_id);
      break;
    case LockMode::EXCLUSIVE:
      txn->GetExclusiveTableLockSet()->insert(table_id);
      break;
  }
}

auto LockManager::LockTable(Transaction *txn, LockMode lock_mode, const table_oid_t &oid) -> bool {
  LockMode held_mode;
  TransactionLevelCheck(txn, lock_mode);
  auto is_upgraded_request = IsTransactionHoldTableLock(txn, oid, &held_mode);
  if (is_upgraded_request) {
    if (lock_mode == held_mode) {
      return true;
    }
    auto upgrade_res = UpgradeLockTable(txn, held_mode, lock_mode, oid);
    return upgrade_res;
  }
  table_lock_map_latch_.lock();
  if (table_lock_map_.find(oid) != table_lock_map_.end()) {
    auto queue_ptr = table_lock_map_[oid];
    table_lock_map_latch_.unlock();
    std::unique_lock<std::mutex> unique_lock{queue_ptr->latch_};
    auto *request = new LockRequest{txn->GetTransactionId(), lock_mode, oid};
    if (queue_ptr->request_queue_.empty() && IsLockModeCompatible(lock_mode, queue_ptr)) {
      request->granted_ = true;
      queue_ptr->granted_queue_.push_back(request);
      InsertLockRecordFromTransaction(lock_mode, txn, oid);
    } else {
      queue_ptr->request_queue_.push_back(request);
      auto res = BlockedUtilRequestGranted(txn, unique_lock, queue_ptr->cv_, request, queue_ptr);
      if (!res) {
        queue_ptr->request_queue_.remove(request);
        delete request;
        GrantNewLocksIfPossible(queue_ptr);
        return false;
      }
      MoveGrantedRequestToGrantedQueue(queue_ptr, request);
      GrantNewLocksIfPossible(queue_ptr);
    }
  } else {
    CreateEntryInTableLockMap(txn, oid, lock_mode);
    table_lock_map_latch_.unlock();
  }
  InsertLockRecordFromTransaction(lock_mode, txn, oid);
  return true;
}

auto LockManager::UpgradeLockTable(Transaction *txn, LockMode origin_mode, LockMode lock_mode, const table_oid_t &oid)
    -> bool {
  IsValidUpgradeRequest(txn, origin_mode, lock_mode);
  table_lock_map_latch_.lock();
  auto queue_ptr = table_lock_map_[oid];
  table_lock_map_latch_.unlock();
  std::unique_lock<std::mutex> unique_lock{queue_ptr->latch_};
  if (queue_ptr->upgrading_ != INVALID_TXN_ID) {
    if (lock_mode == LockMode::EXCLUSIVE || lock_mode == LockMode::SHARED_INTENTION_EXCLUSIVE) {
      txn->SetState(TransactionState::ABORTED);
      throw TransactionAbortException(txn->GetTransactionId(), AbortReason::UPGRADE_CONFLICT);
    }
  }
  queue_ptr->upgrading_ = txn->GetTransactionId();
  std::unordered_set<LockMode> lock_mode_set{};
  GetLockStatesFromQueue(lock_mode_set, queue_ptr, true, txn->GetTransactionId());
  auto upgrade_compatible = CompatibleWithLockModeSet(lock_mode, lock_mode_set);
  auto request_iter = GetExistingLockRequestFromGrantQueue(txn->GetTransactionId(), queue_ptr);
  if (request_iter == queue_ptr->granted_queue_.end()) {
    return false;
  }
  if (upgrade_compatible) {
    // Simply replace it with no lock
    (*request_iter)->lock_mode_ = lock_mode;
    DeleteLockRecordFromTransaction(origin_mode, txn, oid);
    GrantNewLocksIfPossible(queue_ptr);
    InsertLockRecordFromTransaction(lock_mode, txn, oid);
  } else {
    auto *upgrade_request = new LockRequest{txn->GetTransactionId(), lock_mode, oid};
    upgrade_request->upgrade_ = true;
    queue_ptr->request_queue_.push_front(upgrade_request);

    auto *ptr = *request_iter;
    queue_ptr->granted_queue_.erase(request_iter);
    DeleteLockRecordFromTransaction(origin_mode, txn, oid);
    delete ptr;

    auto res = BlockedUtilRequestGranted(txn, unique_lock, queue_ptr->cv_, upgrade_request, queue_ptr);
    if (!res) {
      queue_ptr->request_queue_.remove(upgrade_request);
      delete upgrade_request;
      queue_ptr->upgrading_ = INVALID_TXN_ID;
      GrantNewLocksIfPossible(queue_ptr);
      return false;
    }

    InsertLockRecordFromTransaction(lock_mode, txn, oid);

    MoveGrantedRequestToGrantedQueue(queue_ptr, upgrade_request);

    if (txn->GetTransactionId() == queue_ptr->upgrading_) {
      queue_ptr->upgrading_ = INVALID_TXN_ID;
    }

    GrantNewLocksIfPossible(queue_ptr);
  }
  return true;
}

void LockManager::AppendRequestAfterGrantedRequest(std::shared_ptr<LockRequestQueue> &queue_ptr, LockRequest *request) {
  auto it = queue_ptr->request_queue_.begin();
  while (it != queue_ptr->request_queue_.end()) {
    if ((*it)->granted_) {
      break;
    }
    it++;
  }
  queue_ptr->request_queue_.insert(it, request);
}

void LockManager::GrantLockRequestIfPossible(std::shared_ptr<LockRequestQueue> &queue_ptr, bool upgrade, txn_id_t txn) {
  std::unordered_set<LockMode> lock_state_set{};
  GetLockStatesFromQueue(lock_state_set, queue_ptr, upgrade, txn);
  auto it = queue_ptr->request_queue_.begin();
  if (CompatibleWithLockModeSet((*it)->lock_mode_, lock_state_set)) {
    (*it)->granted_ = true;
    queue_ptr->cv_.notify_all();
  }
}

void LockManager::DeleteLockRequestForUnlock(Transaction *txn, std::shared_ptr<LockRequestQueue> &queue_ptr,
                                             const table_oid_t &oid) {
  auto req_iter = GetExistingLockRequestFromGrantQueue(txn->GetTransactionId(), queue_ptr);
  if (req_iter == queue_ptr->granted_queue_.end()) {
    return;
  }
  auto *ptr = *req_iter;
  DeleteLockRecordFromTransaction(ptr->lock_mode_, txn, oid);
  queue_ptr->granted_queue_.erase(req_iter);
  delete ptr;
}

void LockManager::DeleteLockRequestForUnlockOnRow(Transaction *txn, std::shared_ptr<LockRequestQueue> &queue_ptr,
                                                  const table_oid_t &oid, const RID &rid) {
  auto req_iter = GetExistingLockRequestFromGrantQueue(txn->GetTransactionId(), queue_ptr);
  if (req_iter == queue_ptr->granted_queue_.end()) {
    return;
  }
  auto *ptr = *req_iter;
  DeleteLockRecordFromTransactionForRow(ptr->lock_mode_, txn, oid, rid);
  queue_ptr->granted_queue_.erase(req_iter);
  delete ptr;
}

auto LockManager::UnlockTable(Transaction *txn, const table_oid_t &oid) -> bool {
  // Declared to reuse codes
  HoldNoRowLocksBeforeUnlockTable(txn, oid);
  LockMode param{};
  if (!IsTransactionHoldTableLock(txn, oid, &param)) {
    // Invalid Unlock operation for not holding any lock yet
    txn->SetState(TransactionState::ABORTED);
    throw TransactionAbortException(txn->GetTransactionId(), AbortReason::ATTEMPTED_UNLOCK_BUT_NO_LOCK_HELD);
    return false;
  }
  table_lock_map_latch_.lock();
  if (table_lock_map_.find(oid) == table_lock_map_.end()) {
    table_lock_map_latch_.unlock();
    return false;
  }
  auto queue_ptr = table_lock_map_[oid];
  table_lock_map_latch_.unlock();
  auto unique_lock = std::unique_lock<std::mutex>{queue_ptr->latch_};
  DeleteLockRequestForUnlock(txn, queue_ptr, oid);
  GrantNewLocksIfPossible(queue_ptr);
  ChangeTransactionStateOnUnlock(txn, param);
  return true;
}

auto LockManager::UpgradeLockRow(Transaction *txn, LockMode origin_mode, LockMode lock_mode, const table_oid_t &oid,
                                 const RID &rid) -> bool {
  IsValidUpgradeRequest(txn, origin_mode, lock_mode);
  row_lock_map_latch_.lock();
  auto queue_ptr = row_lock_map_[rid];
  row_lock_map_latch_.unlock();
  std::unique_lock<std::mutex> unique_lock{queue_ptr->latch_};
  if (queue_ptr->upgrading_ != INVALID_TXN_ID) {
    txn->SetState(TransactionState::ABORTED);
    throw TransactionAbortException(txn->GetTransactionId(), AbortReason::UPGRADE_CONFLICT);
  }
  queue_ptr->upgrading_ = txn->GetTransactionId();
  std::unordered_set<LockMode> lock_mode_set{};
  GetLockStatesFromQueue(lock_mode_set, queue_ptr, true, txn->GetTransactionId());
  auto upgrade_compatible = CompatibleWithLockModeSet(lock_mode, lock_mode_set);
  auto request_iter = GetExistingLockRequestFromGrantQueue(txn->GetTransactionId(), queue_ptr);
  if (request_iter == queue_ptr->granted_queue_.end()) {
    return false;
  }
  if (upgrade_compatible) {
    // Simply replace it with no lock
    (*request_iter)->lock_mode_ = lock_mode;
    DeleteLockRecordFromTransactionForRow(origin_mode, txn, oid, rid);
    GrantNewLocksIfPossible(queue_ptr);
    InsertLockRecordFromTransactionForRow(lock_mode, txn, oid, rid);
  } else {
    auto *upgrade_request = new LockRequest{txn->GetTransactionId(), lock_mode, oid, rid};
    upgrade_request->upgrade_ = true;
    queue_ptr->request_queue_.push_front(upgrade_request);

    // Damn bro, we need make a upgrade request
    auto *ptr = *request_iter;
    queue_ptr->granted_queue_.erase(request_iter);
    DeleteLockRecordFromTransactionForRow(origin_mode, txn, oid, rid);
    delete ptr;

    auto res = BlockedUtilRequestGranted(txn, unique_lock, queue_ptr->cv_, upgrade_request, queue_ptr);
    if (!res) {
      queue_ptr->request_queue_.remove(upgrade_request);
      delete upgrade_request;
      queue_ptr->upgrading_ = INVALID_TXN_ID;
      GrantNewLocksIfPossible(queue_ptr);
      return false;
    }

    MoveGrantedRequestToGrantedQueue(queue_ptr, upgrade_request);

    if (txn->GetTransactionId() == queue_ptr->upgrading_) {
      queue_ptr->upgrading_ = INVALID_TXN_ID;
    }

    GrantNewLocksIfPossible(queue_ptr);

    InsertLockRecordFromTransactionForRow(lock_mode, txn, oid, rid);
  }
  return true;
}

auto LockManager::CanTxnTakeLock(Transaction *txn, LockMode lock_mode) -> bool {
  switch (txn->GetIsolationLevel()) {
    case IsolationLevel::REPEATABLE_READ:
      return txn->GetState() == TransactionState::GROWING;
    case IsolationLevel::READ_COMMITTED:
      return txn->GetState() == TransactionState::GROWING ||
             (txn->GetState() == TransactionState::SHRINKING &&
              (lock_mode == LockMode::INTENTION_SHARED || lock_mode == LockMode::SHARED));
    case IsolationLevel::READ_UNCOMMITTED:
      auto lock_exclusive = lock_mode == LockMode::EXCLUSIVE || lock_mode == LockMode::INTENTION_EXCLUSIVE;
      if (!lock_exclusive) {
        txn->SetState(TransactionState::ABORTED);
        throw TransactionAbortException(txn->GetTransactionId(), AbortReason::LOCK_SHARED_ON_READ_UNCOMMITTED);
      }
      return lock_exclusive && txn->GetState() == TransactionState::GROWING;
  }
  return false;
}

void LockManager::GrantNewLocksIfPossible(std::shared_ptr<LockRequestQueue> &lock_request_queue) {
  if (!lock_request_queue->request_queue_.empty()) {
    auto first_entry = lock_request_queue->request_queue_.begin();
    auto *ptr = *first_entry;
    GrantLockRequestIfPossible(lock_request_queue, ptr->upgrade_, ptr->txn_id_);
  }
}

auto LockManager::FindCycle(txn_id_t source_txn, std::vector<txn_id_t> &path, std::unordered_set<txn_id_t> &on_path,
                            std::unordered_set<txn_id_t> &visited, txn_id_t *abort_txn_id) -> bool {
  auto &vec = waits_for_[source_txn];
  if (vec.empty()) {
    return false;
  }
  for (auto txn_id : vec) {
    if (on_path.find(txn_id) != on_path.end()) {
      return false;
    }
    if (visited.find(txn_id) == visited.end()) {
      visited.insert(txn_id);
      path.push_back(txn_id);
      if (FindCycle(txn_id, path, on_path, visited, abort_txn_id)) {
        return true;
      }
      visited.erase(txn_id);
      path.pop_back();
    }
    auto picked_one = PickLowestTxnIdFromPath(path, txn_id);
    *abort_txn_id = picked_one;
    return true;
  }
  on_path.insert(source_txn);
  return false;
}

auto LockManager::PickLowestTxnIdFromPath(std::vector<txn_id_t> &vec, txn_id_t same_idx) -> txn_id_t {
  auto iter = vec.begin();
  while (iter != vec.begin()) {
    if (*iter == same_idx) {
      break;
    }
    iter++;
  }
  txn_id_t target = same_idx;
  while (iter != vec.end()) {
    if (*iter > target) {
      target = *iter;
    }
    iter++;
  }
  return target;
}

inline void LockManager::DeleteLockRecordFromTransactionForRow(LockMode mode, Transaction *txn,
                                                               const table_oid_t &table_id, const RID &row_id) {
  switch (mode) {
    case LockMode::SHARED:
      txn->GetSharedRowLockSet()->at(table_id).erase(row_id);
      break;
    case LockMode::EXCLUSIVE:
      txn->GetExclusiveRowLockSet()->at(table_id).erase(row_id);
      break;
    default:
      // No op
      (void)mode;
  }
}

inline void LockManager::InsertLockRecordFromTransactionForRow(LockMode mode, Transaction *txn,
                                                               const table_oid_t &table_id, const RID &row_id) {
  switch (mode) {
    case LockMode::SHARED:
      if (txn->GetSharedRowLockSet()->find(table_id) == txn->GetSharedRowLockSet()->end()) {
        txn->GetSharedRowLockSet()->insert({table_id, {}});
      }
      txn->GetSharedRowLockSet()->at(table_id).insert(row_id);
      break;
    case LockMode::EXCLUSIVE:
      if (txn->GetExclusiveRowLockSet()->find(table_id) == txn->GetExclusiveRowLockSet()->end()) {
        txn->GetExclusiveRowLockSet()->insert({table_id, {}});
      }
      txn->GetExclusiveRowLockSet()->at(table_id).insert(row_id);
      break;
    default:
      // No op
      (void)mode;
  }
}

void LockManager::ConvertSetToVectorSorted(std::unordered_set<txn_id_t> &txn_id_set, std::vector<txn_id_t> &vec) {
  for (auto txn_id : txn_id_set) {
    vec.push_back(txn_id);
  }
  std::sort(vec.begin(), vec.end());
}

void LockManager::GetHoldersAndWaitersFromResource(std::unordered_set<txn_id_t> &resource_holders,
                                                   std::unordered_set<txn_id_t> &resource_waiters,
                                                   std::shared_ptr<LockRequestQueue> &queue_ptr) {
  std::unique_lock<std::mutex> queue_lock{queue_ptr->latch_};
  auto it = queue_ptr->request_queue_.begin();
  while (it != queue_ptr->request_queue_.end()) {
    if ((*it)->granted_) {
      resource_holders.insert((*it)->txn_id_);
      it++;
    } else {
      break;
    }
  }
  while (it != queue_ptr->request_queue_.end()) {
    if (resource_holders.find((*it)->txn_id_) != resource_holders.end()) {
      it++;
    } else {
      resource_waiters.insert((*it)->txn_id_);
      it++;
    }
  }
}

auto LockManager::LockRow(Transaction *txn, LockMode lock_mode, const table_oid_t &oid, const RID &rid) -> bool {
  CheckLocksAllowedOnRows(txn, lock_mode);
  HoldTableLockBeforeApplyRowLock(txn, lock_mode, oid);
  TransactionLevelCheck(txn, lock_mode);
  LockMode origin_mode;
  if (IsTransactionHoldTableRowLock(txn, oid, rid, &origin_mode)) {
    if (origin_mode == lock_mode) {
      return true;
    }
    auto upgrade_res = UpgradeLockRow(txn, origin_mode, lock_mode, oid, rid);
    return upgrade_res;
  }
  row_lock_map_latch_.lock();
  if (row_lock_map_.find(rid) == row_lock_map_.end()) {
    CreateEntryInTableRowLockMap(txn, oid, rid, lock_mode);
    row_lock_map_latch_.unlock();
  } else {
    auto queue_ptr = row_lock_map_[rid];
    row_lock_map_latch_.unlock();
    std::unique_lock<std::mutex> unique_lock{queue_ptr->latch_};
    if (queue_ptr->request_queue_.empty() && IsLockModeCompatible(lock_mode, queue_ptr)) {
      auto *request = new LockRequest{txn->GetTransactionId(), lock_mode, oid, rid};
      request->granted_ = true;
      queue_ptr->granted_queue_.push_back(request);
    } else {
      auto *request = new LockRequest{txn->GetTransactionId(), lock_mode, oid, rid};
      request->granted_ = false;
      queue_ptr->request_queue_.push_back(request);
      auto res = BlockedUtilRequestGranted(txn, unique_lock, queue_ptr->cv_, request, queue_ptr);
      if (!res) {
        queue_ptr->request_queue_.remove(request);
        delete request;
        GrantNewLocksIfPossible(queue_ptr);
        return false;
      }
      MoveGrantedRequestToGrantedQueue(queue_ptr, request);
      GrantNewLocksIfPossible(queue_ptr);
    }
  }
  InsertLockRecordFromTransactionForRow(lock_mode, txn, oid, rid);
  return true;
}

auto LockManager::UnlockRow(Transaction *txn, const table_oid_t &oid, const RID &rid, bool force) -> bool {
  LockMode param{};
  if (!IsTransactionHoldTableRowLock(txn, oid, rid, &param)) {
    // Invalid Unlock operation for not holding any lock yet
    txn->SetState(TransactionState::ABORTED);
    throw TransactionAbortException(txn->GetTransactionId(), AbortReason::ATTEMPTED_UNLOCK_BUT_NO_LOCK_HELD);
    return false;
  }
  row_lock_map_latch_.lock();
  if (row_lock_map_.find(rid) == row_lock_map_.end()) {
    row_lock_map_latch_.unlock();
    return false;
  }
  auto queue_ptr = row_lock_map_[rid];
  row_lock_map_latch_.unlock();

  std::unique_lock<std::mutex> unique_lock{queue_ptr->latch_};
  DeleteLockRequestForUnlockOnRow(txn, queue_ptr, oid, rid);
  GrantNewLocksIfPossible(queue_ptr);
  if (!force) {
    ChangeTransactionStateOnUnlock(txn, param);
  }
  return true;
}

void LockManager::AddEdge(txn_id_t t1, txn_id_t t2) { waits_for_[t1].insert(t2); }

void LockManager::RemoveEdge(txn_id_t t1, txn_id_t t2) { waits_for_[t1].erase(t2); }

auto LockManager::Dfs(txn_id_t txn_id) -> bool {
  has_search_[txn_id] = true;
  stk_.push_back(txn_id);
  in_stk_[txn_id] = true;

  for (auto x : waits_for_[txn_id]) {
    if (!has_search_[x]) {
      if (Dfs(x)) {
        return true;
      }
    } else if (in_stk_[x]) {
      auto iter = find(stk_.begin(), stk_.end(), x);
      auto max_txn = x;
      while (iter != stk_.end()) {
        if (*iter > max_txn) {
          max_txn = *iter;
        }
        iter++;
      }
      stk_.push_back(max_txn);
      return true;
    }
  }

  stk_.pop_back();
  in_stk_[txn_id] = false;

  return false;
}

auto LockManager::HasCycle(txn_id_t *txn_id) -> bool {
  auto it = waits_for_.begin();
  for (; it != waits_for_.end(); it++) {
    if (!has_search_[it->first] && Dfs(it->first)) {
      *txn_id = stk_.back();
      stk_.pop_back();
      stk_.clear();
      in_stk_.clear();
      has_search_.clear();
      return true;
    }
  }
  stk_.clear();
  in_stk_.clear();
  has_search_.clear();
  return false;
}

auto LockManager::GetEdgeList() -> std::vector<std::pair<txn_id_t, txn_id_t>> {
  std::vector<std::pair<txn_id_t, txn_id_t>> edges(0);
  auto it = waits_for_.begin();
  for (; it != waits_for_.end(); it++) {
    for (auto x : it->second) {
      std::pair<txn_id_t, txn_id_t> pair{it->first, x};
      edges.emplace_back(pair);
    }
  }
  return edges;
}

template <typename K>
void LockManager::BuildGraphFromTable(std::unordered_map<K, std::shared_ptr<LockRequestQueue>> table) {
  for (auto [tid, queue] : table) {
    std::unordered_set<txn_id_t> resource_holders;
    std::unordered_set<txn_id_t> resource_waiters;
    GetHoldersAndWaitersFromResource(resource_holders, resource_waiters, queue);
    std::vector<txn_id_t> waiters_vec{};
    ConvertSetToVectorSorted(resource_waiters, waiters_vec);
    for (auto source : resource_holders) {
      for (auto waiter : waiters_vec) {
        AddEdge(source, waiter);
      }
    }
  }
}

void LockManager::BuildWaitsForGraph() {
  table_lock_map_latch_.lock();
  for (auto &[k, v] : table_lock_map_) {
    v->latch_.lock();
    auto iter1 = v->granted_queue_.begin();
    auto iter2 = v->request_queue_.begin();
    for (; iter1 != v->granted_queue_.end(); iter1++) {
      for (; iter2 != v->request_queue_.end(); iter2++) {
        auto lr1 = *iter1;
        auto lr2 = *iter2;
        if (txn_manager_->GetTransaction(lr1->txn_id_)->GetState() != TransactionState::ABORTED &&
            txn_manager_->GetTransaction(lr2->txn_id_)->GetState() != TransactionState::ABORTED &&
            !Compatible(lr1->lock_mode_, lr2->lock_mode_)) {
          AddEdge(lr1->txn_id_, lr2->txn_id_);
        }
      }
    }
    v->latch_.unlock();
  }
  table_lock_map_latch_.unlock();

  row_lock_map_latch_.lock();
  for (auto &[k, v] : row_lock_map_) {
    v->latch_.lock();
    auto iter1 = v->granted_queue_.begin();
    auto iter2 = v->request_queue_.begin();
    for (; iter1 != v->granted_queue_.end(); iter1++) {
      for (; iter2 != v->request_queue_.end(); iter2++) {
        auto lr1 = *iter1;
        auto lr2 = *iter2;
        if (!Compatible(lr1->lock_mode_, lr2->lock_mode_)) {
          AddEdge(lr1->txn_id_, lr2->txn_id_);
        }
      }
    }
    v->latch_.unlock();
  }
  row_lock_map_latch_.unlock();
}

void LockManager::RemoveAllAboutAbortTxn(txn_id_t abort_id) {
  // 移除请求
  table_lock_map_latch_.lock();
  for (auto &[k, v] : table_lock_map_) {
    v->latch_.lock();

    for (auto iter = v->request_queue_.begin(); iter != v->request_queue_.end();) {
      auto lr = *iter;
      if (lr->txn_id_ == abort_id) {
        v->cv_.notify_all();
      }
      iter++;
    }

    v->latch_.unlock();
  }
  table_lock_map_latch_.unlock();

  row_lock_map_latch_.lock();
  for (auto &[k, v] : row_lock_map_) {
    v->latch_.lock();

    for (auto iter = v->request_queue_.begin(); iter != v->request_queue_.end();) {
      auto lr = *iter;
      if (lr->txn_id_ == abort_id) {
        v->cv_.notify_all();
      }
      iter++;
    }
    v->latch_.unlock();
  }
  row_lock_map_latch_.unlock();

  // 删除边
  waits_for_.erase(abort_id);

  for (auto iter = waits_for_.begin(); iter != waits_for_.end();) {
    if ((*iter).second.count(abort_id) != 0) {
      RemoveEdge((*iter).first, abort_id);
    }
    if ((*iter).second.empty()) {
      waits_for_.erase(iter++);
    } else {
      iter++;
    }
  }
}

void AddSet(std::unordered_set<table_oid_t> &table_set, const std::unordered_set<table_oid_t> &source_set) {
  auto it = source_set.cbegin();
  while (source_set.cend() != it) {
    table_set.insert(*it++);
  }
}

void LockManager::RetriveLockedTableFromTransactions(std::unordered_set<table_oid_t> &table_set, Transaction *txn) {}

void LockManager::RetriveLockedRowFromTransactions(std::unordered_set<RID> &table_set, Transaction *txn) {}

void LockManager::RemoveAllFromQueue(std::shared_ptr<LockRequestQueue> &queue, txn_id_t txn_id) {
  std::unique_lock<std::mutex> unique_lock{queue->latch_};
  auto it = queue->request_queue_.begin();
  while (it != queue->request_queue_.end()) {
    if ((*it)->txn_id_ == txn_id && !(*it)->granted_) {
      auto *ptr = *it;
      queue->request_queue_.erase(it++);
      delete ptr;
    } else {
      it++;
    }
  }
  if (queue->upgrading_ == txn_id) {
    queue->upgrading_ = INVALID_TXN_ID;
  }
  queue->cv_.notify_all();
}

void LockManager::ReleaseAllResourcesHandleByTransaction(txn_id_t txn_id) {
  auto *txn = txn_manager_->GetTransaction(txn_id);
  txn->SetState(TransactionState::ABORTED);
  table_lock_map_latch_.lock();
  auto it = table_lock_map_.begin();
  while (it != table_lock_map_.end()) {
    RemoveAllFromQueue(it->second, txn_id);
    it++;
  }
  table_lock_map_latch_.unlock();
  row_lock_map_latch_.lock();
  auto iter = row_lock_map_.begin();
  while (iter != row_lock_map_.end()) {
    RemoveAllFromQueue(iter->second, txn_id);
    iter++;
  }
  row_lock_map_latch_.unlock();
}

void LockManager::RunCycleDetection() {
  while (enable_cycle_detection_) {
    std::this_thread::sleep_for(cycle_detection_interval);
    {  // TODO(students): detect deadlock
      BuildWaitsForGraph();
      while (true) {
        txn_id_t abort_tid;
        if (HasCycle(&abort_tid)) {
          txn_manager_->GetTransaction(abort_tid)->SetState(TransactionState::ABORTED);
          RemoveAllAboutAbortTxn(abort_tid);
        } else {
          break;
        }
      }
      waits_for_.clear();
    }
  }
}

void LockManager::UnlockAll() {
  table_lock_map_latch_.lock();
  for (auto [t_id, queue_ptr] : table_lock_map_) {
    auto it = queue_ptr->request_queue_.begin();
    while (it != queue_ptr->request_queue_.end()) {
      auto *ptr = *it;
      queue_ptr->request_queue_.erase(it++);
      delete ptr;
    }
    auto it2 = queue_ptr->granted_queue_.begin();
    while (it2 != queue_ptr->granted_queue_.end()) {
      auto *ptr = *it2;
      queue_ptr->granted_queue_.erase(it2++);
      delete ptr;
    }
  }
  table_lock_map_latch_.unlock();
  row_lock_map_latch_.lock();
  for (auto [t_id, queue_ptr] : row_lock_map_) {
    auto it = queue_ptr->request_queue_.begin();
    while (it != queue_ptr->request_queue_.end()) {
      auto *ptr = *it;
      queue_ptr->request_queue_.erase(it++);
      delete ptr;
    }
    auto it2 = queue_ptr->granted_queue_.begin();
    while (it2 != queue_ptr->granted_queue_.end()) {
      auto *ptr = *it2;
      queue_ptr->granted_queue_.erase(it2++);
      delete ptr;
    }
  }
  row_lock_map_latch_.unlock();
}

}  // namespace bustub
